package com.springai.demo.controller;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetriever;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import com.alibaba.cloud.ai.toolcalling.amp.WeatherSearchService;
import com.springai.demo.record.ActorsFilms;
import com.springai.demo.tools.TimeFunction;
import com.springai.demo.tools.TimeTools;
import com.springai.demo.tools.UserInfoTools;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.image.*;
import org.springframework.ai.rag.Query;
import org.springframework.ai.tool.function.FunctionToolCallback;
import org.springframework.ai.tool.metadata.ToolMetadata;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * @author hh
 * @since 2025/7/15
 */
@RestController
public class ChatController {

    @Autowired
    private DashScopeChatModel dashScopeChatModel;

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ImageModel imageModel;

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private WeatherSearchService weatherSearchService;

    @GetMapping("chatModel")
    public Flux<String> chat(String message) {
        Flux<ChatResponse> stream = dashScopeChatModel.stream(new Prompt(message, DashScopeChatOptions.builder()
                .withTopP(0.7)
                .withTopK(50)
                .withTemperature(0.8)
                .withModel(DashScopeApi.ChatModel.QWEN_PLUS.getModel())
                .build()));
        return stream.map(resp -> resp.getResult().getOutput().getText());
        // return dashScopeChatModel.stream(message);
    }

    /*
        Top-K：将 token 的选择范围限制为 K 个最有可能的后续 token。值越高（例如 40-50），多样性越好。
        Top-P（核采样）：从最小的标记集中动态选择，其累积概率超过P。0.8-0.95 之类的值很常见。
     */
    @GetMapping("chatClient")
    public Flux<String> chatClient(String message, String id) {
        return chatClient.prompt()
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, id))
                .system("你是一个天气助手，只回答有关天气的内容，非天气相关的只回答：无可奉告")
                .user(message)
                .options(DashScopeChatOptions.builder()
                        .withModel(DashScopeApi.ChatModel.QWEN_MAX.getModel())
                        .withTemperature(0.4)
                        .withTopK(40)
                        .withTopP(0.8)
                        .withMaxToken(200)
                        .build())
                .stream()
                .content();
    }

    /**
     * 基于Mysql的对话记忆
     */
    @GetMapping("memory")
    public Flux<String> memory(String message, String id) {
        return chatClient.prompt(message)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, id))
                .advisors(new SimpleLoggerAdvisor())
                .stream()
                .content();
    }

    /**
     * 基于内存的对话记忆
     */
    @GetMapping("inMemory")
    public String inMemory() {
        ChatMemoryRepository chatMemoryRepository = new InMemoryChatMemoryRepository();
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .build();
        ChatClient client = ChatClient.builder(dashScopeChatModel)
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                .build();

        // 对话记忆的唯一标识
        String conversantId = UUID.randomUUID().toString();

        ChatResponse chatResponse = client.prompt()
                .user("我想去新疆")
                .advisors(spec -> spec.param("conversantId", conversantId)
                        .param("retriveSize", 10))
                .call()
                .chatResponse();
        String content = chatResponse.getResult().getOutput().getText();
        System.out.println(content);

        chatResponse = client.prompt()
                .user("可以帮我推荐一些美食吗？")
                .advisors(spec -> spec.param("conversantId", conversantId)
                        .param("retriveSize", 10))
                .call()
                .chatResponse();
        content = chatResponse.getResult().getOutput().getText();
        System.out.println(content);
        return "ok";
    }

    @GetMapping("image")
    public String image(String message) {
        ImageOptions imageOptions = ImageOptionsBuilder.builder()
                .height(800)
                .width(1024)
                .model("wanx2.0-t2i-turbo")
                .build();
        ImagePrompt imagePrompt = new ImagePrompt(message, imageOptions);
        ImageResponse imageResponse = imageModel.call(imagePrompt);
        return imageResponse.getResult().getOutput().getUrl();
    }

    @GetMapping("embedding")
    public Map embed(String message) {
        EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(message));
        return Map.of("embedding", embeddingResponse);
    }

    @GetMapping("multiModal")
    public Flux<String> multiModal() {
        ClassPathResource classPathResource = new ClassPathResource("cat.png");
        dashScopeChatModel.setDashScopeChatOptions(DashScopeChatOptions.builder()
                .withModel("qwen-vl-max")
                .withTemperature(0.7)
                .build());
        return ChatClient.create(dashScopeChatModel).prompt()
                .user(u -> u.text("描述一下图片的内容")
                        .media(MimeTypeUtils.IMAGE_PNG, classPathResource))
                .stream()
                .content();
    }

    @GetMapping("template")
    public Flux<String> template(String city) {
        String messageTemplate = "介绍一下{city}的旅游景点";
        PromptTemplate promptTemplate = new PromptTemplate(messageTemplate);
        Prompt prompt = promptTemplate.create(Map.of("city", city));
        Flux<ChatResponse> stream = dashScopeChatModel.stream(prompt);
        return stream.map(t -> t.getResult().getOutput().getText());
    }

    /*
        该示例将创建一个DashScopeDocumentRetriever实例，在指定的知识库上执行检索操作。
     */
    @GetMapping("retrieve")
    public void retrieve(String message) {
        DashScopeApi dashScopeApi = DashScopeApi.builder()
                .apiKey("xxxxxx")
                .build();
        DashScopeDocumentRetriever dashScopeDocumentRetriever = new DashScopeDocumentRetriever(dashScopeApi,
                DashScopeDocumentRetrieverOptions.builder()
                        .withIndexName("1")
                        .build());
        List<Document> retrieve = dashScopeDocumentRetriever.retrieve(Query.builder()
                .text(message)
                .build());
        for (Document document : retrieve) {
            System.out.println(document.getText());
        }
    }

    @GetMapping("structure")
    public String structure(String message) {
        ActorsFilms actorsFilms = ChatClient.create(dashScopeChatModel)
                .prompt()
                .user(u -> u.text("输出{actor}演的5个电影.")
                        .param("actor", message))
                .call()
                .entity(ActorsFilms.class);
        return actorsFilms.toString();
    }

    @GetMapping("structureList")
    public String structureList() {
        List<ActorsFilms> actorsFilmsList = ChatClient.create(dashScopeChatModel)
                .prompt()
                .user("Generate the filmography of 5 movies for Tom Hanks and Bill Murray.")
                .call()
                .entity(new ParameterizedTypeReference<List<ActorsFilms>>() {
                });
        return actorsFilmsList.toString();
    }

    /*
        函数调用不支持
        Optional
        异步类型（CompletableFuture、Future）
        响应式类型（Flow、Mono、Flux）
        函数类型（Function、Supplier、Consumer）
     */
    @GetMapping("tools")
    public String tools(String message, String id) {
        return chatClient.prompt(message)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, id))
                .tools(new TimeTools())
                .call()
                .content();
    }

    @GetMapping("function")
    public String function(String message, String id) {
        return chatClient.prompt(message)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, id))
                .toolCallbacks(FunctionToolCallback
                        .builder("getTimeByZoneId", new TimeFunction())
                        .description("Get time by zone id")
                        /*
                            默认情况下，工具调用的返回值会再次回传到 AI 模型进一步处理。但在一些场景中需要将结果直接返回给调用方而非模型，比如数据搜索。
                            定义方法工具时，可以通过 @Tool 注解的 returnDirect 参数置 true 来启动直接返回；定义方法工具和函数工具时需要通过 ToolMetadata 对象传递到 MethodToolCallBack.Builder 和 FunctionToolCallBack.Builder中。
                            调用这段代码将直接返回 TimeFunction 返回的JSON对象，而不再经过大模型加工处理。
                         */
                        .toolMetadata(ToolMetadata.builder()
                                .returnDirect(true)
                                .build())
                        .inputType(TimeFunction.Request.class)
                        .build())
                .call()
                .content();
    }

    @GetMapping("context")
    public String getUserName(String id) {
        return chatClient.prompt("获取我的用户名")
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, id))
                .tools(new UserInfoTools())
                .toolContext(Map.of("userId", "123"))
                .call()
                .content();
    }

    @GetMapping("vectorStore")
    public void vectorStore() {
        List<Document> documents = new ArrayList<>();
        ClassPathResource resource = new ClassPathResource("张三简历.txt");
        try (BufferedInputStream bis = new BufferedInputStream(resource.getInputStream())){
            byte[] bytes = new byte[1024];
            while (bis.read(bytes) != -1) {
                String s = new String(bytes);
                Document document = new Document(s);
                documents.add(document);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        VectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        vectorStore.add(documents);

        List<Document> results = vectorStore.similaritySearch(SearchRequest.builder()
                .query("张三毕业于哪个大学")
                .topK(2)
                .build());
        for (Document result : results) {
            System.out.println(result.getText());
        }
    }

    @GetMapping("gaodeMap")
    public Flux<String> gaodeMap(String message) {
        return chatClient.prompt()
                .user(message)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "1"))
                .toolCallbacks(FunctionToolCallback
                        .builder("gaoDeGetAddressWeather", weatherSearchService)
                        .description("获取城市天气")
                        .inputType(WeatherSearchService.Request.class)
                        .build())
                // .toolNames("gaoDeGetAddressWeather")
                .stream()
                .content();
    }

}
